import tensorflow as tf

FLAGS = tf.flags.FLAGS

tf.flags.DEFINE_string("checkpoint_path", "",
                       "Model checkpoint file or directory containing a "
                       "model checkpoint file.")
def RenameCkpt(oldPath, newPath):
    vars_to_rename = {
    "lstm/BasicLSTMCell/Linear/Matrix": "lstm/basic_lstm_cell/kernel",
    "lstm/BasicLSTMCell/Linear/Bias": "lstm/basic_lstm_cell/bias",
    }
    new_checkpoint_vars = {}
    reader = tf.train.NewCheckpointReader(oldPath)
    for old_name in reader.get_variable_to_shape_map():
      print(">>>%s" %old_name)
      if old_name in vars_to_rename:
        new_name = vars_to_rename[old_name]
      else:
        new_name = old_name
      new_checkpoint_vars[new_name] = tf.Variable(reader.get_tensor(old_name))

    init = tf.global_variables_initializer()
    saver = tf.train.Saver(new_checkpoint_vars)

    with tf.Session() as sess:
      sess.run(init)
      saver.save(sess, newPath)
    print("checkpoint file rename successful... ")

if __name__ == '__main__':
    oldPath = '/root/data/qinchao/im2txt/im2txt/model/train/model.ckpt-2000000'
    newPath = '/root/data/qinchao/im2txt/im2txt/model/newmodel.ckpt-2000000'
    RenameCkpt(oldPath, newPath)
